from faker import Faker
import pymysql
import openpyxl
import random


class DButils:
    def __init__(self, host, user, password, database, port=3306, charset='utf8'):
        self.host = host
        self.user = user
        self.password = password
        self.database = database
        self.port = port
        self.charset = charset

    def execute(self, sql, data=None):
        """通用执行方法，用于执行INSERT/UPDATE/DELETE/CREATE等语句"""
        db = pymysql.connect(
            host=self.host,
            user=self.user,
            password=self.password,
            database=self.database,
            port=self.port,
            charset=self.charset
        )
        try:
            with db.cursor() as cursor:
                cursor.execute(sql, data)
                db.commit()
                return True
        except Exception as e:
            db.rollback()
            print(f"执行SQL失败: {e}")
            return False
        finally:
            db.close()

    def select(self, sql, data=None):
        """查询数据"""
        db = pymysql.connect(
            host=self.host,
            user=self.user,
            password=self.password,
            database=self.database,
            port=self.port,
            charset=self.charset
        )
        try:
            with db.cursor() as cursor:
                cursor.execute(sql, data)
                return cursor.fetchall()
        except Exception as e:
            print(f'数据查询失败: {e}')
            return None
        finally:
            db.close()

    def batch_insert(self, sql, data_list):
        """批量插入数据"""
        db = pymysql.connect(
            host=self.host,
            user=self.user,
            password=self.password,
            database=self.database,
            port=self.port,
            charset=self.charset
        )
        try:
            with db.cursor() as cursor:
                # 批量插入，提高效率
                cursor.executemany(sql, data_list)
                db.commit()
                print(f"成功插入 {len(data_list)} 条数据")
                return True
        except Exception as e:
            db.rollback()
            print(f"批量插入失败: {e}")
            return False
        finally:
            db.close()


def generate_excel_data(filename, rows=2000):
    """生成指定行数的用户信息Excel数据"""
    fake = Faker('zh_CN')

    # 创建工作簿和工作表
    wb = openpyxl.Workbook()
    ws = wb.active
    ws.title = "用户信息"

    # 定义字段
    fields = [
        "id", "idcard", "username", "realname", "pwd",
        "telphone", "email", "age", "sex", "address",
        "hiredate", "sal", "job", "company"
    ]

    # 写入表头
    for col, field in enumerate(fields, 1):
        ws.cell(row=1, column=col, value=field)

    # 生成并写入数据
    for row in range(2, rows + 2):
        data = [
            row - 1,  # id
            fake.ssn(),  # 身份证号
            fake.user_name(),  # 用户名
            fake.name(),  # 真实姓名
            fake.password(length=8),  # 密码
            fake.phone_number(),  # 电话号码
            fake.email(),  # 邮箱
            random.randint(18, 65),  # 年龄
            random.choice(["男", "女"]),  # 性别
            fake.address().replace("\n", ""),  # 地址
            fake.date_between(start_date='-10y', end_date='today'),  # 入职日期
            round(random.uniform(3000, 20000), 2),  # 薪资
            fake.job(),  # 职位
            fake.company()  # 公司
        ]

        for col, value in enumerate(data, 1):
            ws.cell(row=row, column=col, value=value)

    wb.save(filename)
    print(f"已生成Excel文件: {filename}，共{rows}行数据")
    return filename


def import_excel_to_db(db, excel_file):
    """将Excel数据导入到数据库"""
    try:
        # 打开Excel文件
        wb = openpyxl.load_workbook(excel_file)
        ws = wb.active  # 获取活动工作表

        # 获取数据行数
        row_count = ws.max_row
        print(f"开始导入Excel数据，共 {row_count - 1} 条记录...")

        # 准备插入的数据列表
        data_list = []
        for row in range(2, row_count + 1):  # 从第2行开始（跳过表头）
            data = (
                ws.cell(row=row, column=1).value,  # id
                ws.cell(row=row, column=2).value,  # idcard
                ws.cell(row=row, column=3).value,  # username
                ws.cell(row=row, column=4).value,  # realname
                ws.cell(row=row, column=5).value,  # pwd
                ws.cell(row=row, column=6).value,  # telphone
                ws.cell(row=row, column=7).value,  # email
                ws.cell(row=row, column=8).value,  # age
                ws.cell(row=row, column=9).value,  # sex
                ws.cell(row=row, column=10).value,  # address
                ws.cell(row=row, column=11).value,  # hiredate
                ws.cell(row=row, column=12).value,  # sal
                ws.cell(row=row, column=13).value,  # job
                ws.cell(row=row, column=14).value  # company
            )
            data_list.append(data)

        # 插入数据的SQL语句
        insert_sql = """
        INSERT INTO user_info (
            id, idcard, username, realname, pwd, telphone, email, 
            age, sex, address, hiredate, sal, job, company
        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """

        # 执行批量插入
        return db.batch_insert(insert_sql, data_list)

    except Exception as e:
        print(f"导入Excel到数据库失败: {e}")
        return False


if __name__ == '__main__':
    # 数据库配置，请根据实际情况修改
    db_config = {
        'host': 'localhost',
        'user': 'root',
        'password': '123456',
        'database': 'hrsx',
        'port': 3306
    }

    # 初始化数据库连接
    db = DButils(**db_config)

    # 1. 创建数据表（如果不存在）
    # create_table_sql = """
    # CREATE TABLE IF NOT EXISTS user_info (
    #     id int(11) NOT NULL,
    #     idcard varchar(50) NOT NULL,
    #     username varchar(50) NOT NULL,
    #     realname varchar(50) NOT NULL,
    #     pwd varchar(50) NOT NULL,
    #     telphone varchar(12) NOT NULL,
    #     email varchar(100) NOT NULL,
    #     age int(11) NOT NULL,
    #     sex varchar(20) NOT NULL,
    #     address varchar(200) NOT NULL,
    #     hiredate date NOT NULL,
    #     sal double(9, 2) NOT NULL,
    #     job varchar(100) NOT NULL,
    #     company varchar(100) NOT NULL,
    #     PRIMARY KEY (id)
    # ) ENGINE=InnoDB DEFAULT CHARSET=utf8;
    # """
    # if db.execute(create_table_sql):
    #     print("数据表准备就绪")

    # 2. 生成Excel数据
    excel_file = generate_excel_data("user_info.xlsx", rows=2000)

    # 3. 将Excel数据导入到数据库
    if import_excel_to_db(db, excel_file):
        print("数据导入完成！")

        # 4. 验证数据（查询前10条记录）
        result = db.select("SELECT id, username, realname FROM user_info LIMIT 10")
        if result:
            print("\n数据库中前10条记录:")
            for row in result:
                print(row)
